﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading;
using Drifting.MotionPlanning.Primitives;
using Drifting.MotionPlanning.Systems;
using Drifting.MotionPlanning.PointSets;
using Drifting.Foundation;

namespace Drifting.MotionPlanning.SingleThreading.Algorithms
{
    public class RRTstar : GraphBasedPlanningAlgorithm
    {
        public RRTstar(BaseSystem system, Random rand)
            : base(system, rand)
        {
            int[] branches = new int[] { 8, 8, 8, 8, 8, 8 };   //{ 6, 6, 6, 6, 6, 6, 6, 6, 6 }
            PointSet Vactive = new BinQuery(system.Min_Alias, system.Max_Alias, system.DOF_Alias, branches);

            vertices =
                //new LinearQuery();
                //new AlphaRootQuery(rand, 1.5);
                Vactive;
        }

        override public void Execute(object input)
        {
            if (input is Int32 == false)
                throw new ArgumentException();
            n_exec = (Int32)input;
            if (x_root == null)
                throw new InvalidOperationException("Forget to call SetRoot first?");

			List<Edge> E_near = new List<Edge>();
            double[] a_sample = null;
            n_current = 0;
            while (n_current < n_exec)
            {
                n_current++;

                //Sample
                system.SampleAlias(rand, ref a_sample);
                if (system.inStateSpace(a_sample) == false) continue;
                double[] q_new = new double[system.DOF_State];
                MathAssistant.vector_set(a_sample, ref q_new);
                GraphState x_new = (GraphState)SpawnState(q_new);
                
                List<State> X_near = new List<State>();
                int K = (int)((1 + (int)Math.Ceiling(Math.Log(vertices.Count))) * 1.5);
                vertices.Neighbor_KQuery(a_sample, K, ref X_near);

                //Extend the graph
                E_near.Clear();
                double cost_MIN = double.MaxValue;
                GraphState x_select = null;
                Edge e_select = null;
				for(int i=0;i<X_near.Count;++i)
				{
					GraphState x_near = (GraphState)X_near[i];
                    Edge e_near = system.Steer(x_near, x_new);
                    bool validEdge = system.ValidateEdge(e_near);
                    if (validEdge == false) continue;
                    E_near.Add(e_near);

                    double cost_near = x_near.Cost + e_near.step_cost;
                    if (cost_MIN > cost_near)
                    {
                        x_select = x_near;
                        e_select = e_near;
                    }
				}
                if (e_select == null) continue;

                e_select.entry = x_select;
                e_select.exit = x_new;
                x_new.ParentEdge = e_select;
                x_select.Children.Add(e_select);
                x_new.Cost = x_select.Cost + e_select.step_cost;
                vertices.AddState(x_new);


                
                //Rewire the graph
                E_near.Clear();
                for (int i = 0; i < X_near.Count; ++i)
                {
                    GraphState x_near = (GraphState)X_near[i];
                    Edge e_near = system.Steer(x_near, x_new);
                    bool validEdge = system.ValidateEdge(e_near);
                    if (validEdge == false) continue;
                    E_near.Add(e_near);

                    double cost_near = x_new.Cost + e_near.step_cost;
                    if (cost_near < x_near.Cost)
                    {
                        Edge e_near_parent = x_near.ParentEdge;
                        GraphState x_near_parent = (GraphState)e_near_parent.entry;
                        //TODO null check

                        x_near_parent.Children.Remove(e_near_parent);
                        x_near.ParentEdge = null;

                        e_near.entry = x_new;
                        e_near.exit = x_near;

                        x_new.Children.Add(e_near);
                        x_near.ParentEdge = e_near;
                        x_near.Cost = cost_near;

                        //propagate costs down
                        CostPropagate(x_near);
                    }
                }

                //Optional
                if (inXgoal(q_new))
                {
                    if (cost_solution >= x_new.Cost)
                    {
                        cost_solution = x_new.Cost;
                        x_goal = x_new;
                    }
                }
            }
        }

        private void CostPropagate(GraphState x)
        {
            if (x.Children == null) return;

            double cost_x = x.Cost;
            for (int i = 0; i < x.Children.Count; ++i)
            {
                Edge e_i = x.Children[i];
                double cost_i_new = cost_x + e_i.step_cost;
                GraphState x_i = (GraphState)e_i.exit;
                x_i.Cost = cost_i_new;
                CostPropagate(x_i);
            }
        }

        override public void AliasOf(double[] q, ref double[] alias)
        {
            if (alias == null)
                alias = new double[system.DOF_Alias];
            for (int i = 0; i < system.DOF_Alias; ++i)
                alias[i] = q[i];
        }

    }
}
